import os
import yaml
import random
import shutil
from tqdm import tqdm
import xml.etree.ElementTree as ET

# 创建YOLO需要的数据集格式
def convert_annotations(annotations_dir, images_dir, output_dir, class_names):
    """
    将XML标注转换为YOLO格式
    """
    os.makedirs(output_dir, exist_ok=True)
    
    # 创建类别映射
    class_to_idx = {name: i for i, name in enumerate(class_names)}
    
    xml_files = [f for f in os.listdir(annotations_dir) if f.endswith('.xml')]
    for xml_file in tqdm(xml_files, desc="Converting annotations"):
        # 解析XML文件
        tree = ET.parse(os.path.join(annotations_dir, xml_file))
        root = tree.getroot()
        
        image_filename = root.find('filename').text
        image_width = int(root.find('size/width').text)
        image_height = int(root.find('size/height').text)
        
        # 创建YOLO格式的标注文件（.txt）
        txt_filename = os.path.splitext(image_filename)[0] + '.txt'
        txt_path = os.path.join(output_dir, txt_filename)
        
        with open(txt_path, 'w') as f:
            for obj in root.findall('object'):
                # 获取类别名称和对应的索引
                class_name = obj.find('name').text  # 使用<name>标签
                if class_name not in class_to_idx:
                    continue
                class_idx = class_to_idx[class_name]
                
                # 获取边界框坐标
                xmin = float(obj.find('bndbox/xmin').text)
                ymin = float(obj.find('bndbox/ymin').text)
                xmax = float(obj.find('bndbox/xmax').text)
                ymax = float(obj.find('bndbox/ymax').text)
                
                # 转换为YOLO格式：center_x, center_y, width, height (相对值)
                center_x = (xmin + xmax) / 2.0 / image_width
                center_y = (ymin + ymax) / 2.0 / image_height
                width = (xmax - xmin) / image_width
                height = (ymax - ymin) / image_height
                
                # 写入文件
                f.write(f"{class_idx} {center_x} {center_y} {width} {height}\n")
    
    return len(xml_files)

# 创建数据集配置文件
def create_dataset_config(output_dir, class_names, train_images, val_images, test_images=None):
    """
    创建YOLO所需的数据集配置文件
    """
    data_yaml = {
        'path': os.path.abspath(output_dir),  # 数据集根目录
        'train': train_images,  # 训练集图像路径
        'val': val_images,      # 验证集图像路径
        'nc': len(class_names), # 类别数量
        'names': class_names    # 类别名称列表
    }
    
    if test_images:
        data_yaml['test'] = test_images  # 测试集图像路径
    
    # 保存为YAML文件
    config_path = os.path.join(output_dir, 'dataset.yaml')
    with open(config_path, 'w') as f:
        yaml.dump(data_yaml, f, default_flow_style=False)
    
    return config_path

# 分割数据集
def split_dataset(annotations_dir, output_dir, train_ratio=0.8, val_ratio=0.1, test_ratio=0.1):
    """
    将数据集分割为训练集、验证集和测试集
    """
    os.makedirs(os.path.join(output_dir, 'train', 'labels'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'val', 'labels'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'test', 'labels'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'train', 'images'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'val', 'images'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'test', 'images'), exist_ok=True)
    
    xml_files = [f for f in os.listdir(annotations_dir) if f.endswith('.xml')]
    random.shuffle(xml_files)
    
    n_train = int(len(xml_files) * train_ratio)
    n_val = int(len(xml_files) * val_ratio)
    
    train_files = xml_files[:n_train]
    val_files = xml_files[n_train:n_train+n_val]
    test_files = xml_files[n_train+n_val:]
    
    return train_files, val_files, test_files

# 复制和组织数据集
def organize_dataset(images_dir, annotations_dir, output_dir, train_files, val_files, test_files):
    """
    根据分割结果，复制图像和标注文件到相应目录
    """
    
    # 处理训练集
    for xml_file in tqdm(train_files, desc="Organizing training set"):
        base_name = os.path.splitext(xml_file)[0]
        
        # 复制图像
        src_img = os.path.join(images_dir, f"{base_name}.png")
        dst_img = os.path.join(output_dir, 'train', 'images', f"{base_name}.png")
        if os.path.exists(src_img):
            shutil.copy(src_img, dst_img)
        
        # 复制标注
        src_txt = os.path.join(output_dir, f"{base_name}.txt")
        dst_txt = os.path.join(output_dir, 'train', 'labels', f"{base_name}.txt")
        if os.path.exists(src_txt):
            shutil.copy(src_txt, dst_txt)
    
    # 处理验证集
    for xml_file in tqdm(val_files, desc="Organizing validation set"):
        base_name = os.path.splitext(xml_file)[0]
        
        # 复制图像
        src_img = os.path.join(images_dir, f"{base_name}.png")
        dst_img = os.path.join(output_dir, 'val', 'images', f"{base_name}.png")
        if os.path.exists(src_img):
            shutil.copy(src_img, dst_img)
        
        # 复制标注
        src_txt = os.path.join(output_dir, f"{base_name}.txt")
        dst_txt = os.path.join(output_dir, 'val', 'labels', f"{base_name}.txt")
        if os.path.exists(src_txt):
            shutil.copy(src_txt, dst_txt)
    
    # 处理测试集
    for xml_file in tqdm(test_files, desc="Organizing test set"):
        base_name = os.path.splitext(xml_file)[0]
        
        # 复制图像
        src_img = os.path.join(images_dir, f"{base_name}.png")
        dst_img = os.path.join(output_dir, 'test', 'images', f"{base_name}.png")
        if os.path.exists(src_img):
            shutil.copy(src_img, dst_img)
        
        # 复制标注
        src_txt = os.path.join(output_dir, f"{base_name}.txt")
        dst_txt = os.path.join(output_dir, 'test', 'labels', f"{base_name}.txt")
        if os.path.exists(src_txt):
            shutil.copy(src_txt, dst_txt)

# 准备数据集函数 - 封装数据处理流程
def prepare_dataset(annotations_dir, images_dir, output_dir, class_names):
    """
    准备完整的YOLO格式数据集
    如果YOLO格式数据已经存在，跳过数据转换步骤
    """
    dataset_yaml_path = os.path.join(output_dir, 'dataset.yaml')
    
    # 检查数据集配置文件是否已存在
    if os.path.exists(dataset_yaml_path):
        print(f"发现已存在的数据集配置：{dataset_yaml_path}")
        with open(dataset_yaml_path, 'r') as f:
            data_config = yaml.safe_load(f)
        
        # 检查训练和验证数据目录是否存在
        train_dir = os.path.join(output_dir, 'train')
        val_dir = os.path.join(output_dir, 'val')
        test_dir = os.path.join(output_dir, 'test')
        
        if (os.path.exists(train_dir) and os.path.exists(val_dir) and
            os.path.exists(os.path.join(train_dir, 'labels')) and
            os.path.exists(os.path.join(val_dir, 'labels'))):
            print("发现已准备好的YOLO格式数据，跳过数据处理步骤")
            return dataset_yaml_path
    
    print("未发现完整的已准备数据，开始数据处理流程...")
    
    # 1. 将XML标注转换为YOLO格式
    print("Converting annotations to YOLO format...")
    convert_annotations(annotations_dir, images_dir, output_dir, class_names)
    
    # 2. 分割数据集
    print("Splitting dataset...")
    train_files, val_files, test_files = split_dataset(annotations_dir, output_dir)
    
    # 3. 组织数据集目录结构
    print("Organizing dataset directory structure...")
    organize_dataset(images_dir, annotations_dir, output_dir, train_files, val_files, test_files)
    
    # 4. 创建数据集配置文件
    print("Creating dataset configuration...")
    data_cfg = create_dataset_config(
        output_dir,
        class_names,
        os.path.join(output_dir, 'train', 'images'),
        os.path.join(output_dir, 'val', 'images'),
        os.path.join(output_dir, 'test', 'images')
    )
    
    return data_cfg 